##########################################################################################

library(data.table)
library(optparse)
library(ArchR)
library(ggsci)
library(Seurat)

##########################################################################################
option_list <- list(
    make_option(c("--comine_data_file"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){
    
    ## 整合atac和rna的文件
    comine_data_file <- "~/20231121_singleMuti/results/qc_atac_v3/germ/testis_combined_peak.combineRNA.qc.Rdata"

    ## 输出
    out_path <- "~/20231121_singleMuti/results/trajectory"

}


###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

comine_data_file <- opt$comine_data_file
out_path <- opt$out_path

dir.create(out_path , recursive = T)
image_path <- out_path

###########################################################################################
## 细胞颜色
use_colors <- c(pal_npg("nrc")(10) , pal_jco("default")(6))
#names(use_colors) <- unique(scrnat$cell_type)
names(use_colors) <-c( "Myoid cells","Leydig cells" ,          
"Endothelial cells","Zygotene",         
"Round&ElongateS.tids","Patchytene",
"SSC","Sperm" ,                 
"Diplotene","Early stage of spermatids", 
"Leptotene","Sertoli cells",            
"Macrophages","Differenting&Differented SPG",
"Pericytes","NKT cells")

###########################################################################################

a <- load(comine_data_file)
## testis_combined_peak_combineRNA
projHeme5 <- testis_combined_peak_combineRNA

###########################################################################################
## 演化方向
cell_level <- c("SSC","Differenting&Differented SPG",
    "Leptotene","Zygotene","Patchytene","Diplotene",
    "Early stage of spermatids","Round&ElongateS.tids","Sperm"
    )

## 15-1、1-5、2/10-3
## 有丝分裂、有丝分裂-减数分裂、减数分裂-变态反应
cell_level1 <- c("SSC","Differenting&Differented SPG")
cell_level2 <- c("Differenting&Differented SPG" , "Leptotene")
cell_level3 <- c("Patchytene","Diplotene","Early stage of spermatids")

## 四个方向
cell_level_list <- list(
    all = cell_level ,
    mitosis = cell_level1 ,
    mitosis_meiosis = cell_level2 ,
    meiosis_allergy = cell_level3
    )

## 去除基因
del_gene <- "FOXO1"

###########################################################################################
## 按照不同的时序分别去跑
## https://www.archrproject.com/bookdown/myeloid-trajectory-monocyte-differentiation.html

for(level_name in names(cell_level_list)){

    ## 定义时序关系
    trajectory <- cell_level_list[[level_name]]

    ## 输出在特定的时序的文件夹
    out_path <- paste0( image_path , "/" , level_name )
    dir.create(out_path , recursive = T)

    ###########################################################################################
    ## atac的聚类图,标记细胞类型
    p <- plotEmbedding(ArchRProj = projHeme5, colorBy = "cellColData", name = "cell_type", embedding = "UMAP" , pal = use_colors)
    p1 <- plotEmbedding(ArchRProj = projHeme5, colorBy = "cellColData", name = "Clusters", embedding = "UMAP")
    image_name <- paste0( out_path , "/plotEmbedding.atac.pdf")
    pdf(image_name , width = 10, height = 8)
    print(p)
    print(p1)
    dev.off()

    ###########################################################################################
    ## 计算轨迹
    ## 发现每个单元格都有一个介于 0 到 100 之间的唯一伪时间值。我们排除具有NA值的单元格，因为这些单元格不是轨迹的一部分
    projHeme5 <- addTrajectory(
        ArchRProj = projHeme5, 
        name = "Trajectory_time", 
        groupBy = "cell_type",
        trajectory = trajectory, 
        embedding = "UMAP", 
        force = TRUE
    )

    ## 绘制轨迹，我们使用将plotTrajectory()伪时间值叠加在 UMAP 嵌入上的函数，并显示一个近似样条拟合轨迹路径的箭头。
    ## 不属于轨迹的单元格呈灰色。在此示例中，我们用来colorBy = "cellColData"告诉 ArchR 在内部查找由“MyeloidU”伪时间轨迹cellColData指定的列。
    p <- plotTrajectory(projHeme5, trajectory = "Trajectory_time", colorBy = "cellColData", name = "Trajectory_time")
    image_name <- paste0( out_path , "/plotTrajectory.atac.pdf")
    pdf(image_name , width = 10, height = 8)
    print(p)
    dev.off()

    ## 保存计算完时间的文件
    out_file <- paste0( out_path , "/plotTrajectory.atac.Rdata")
    save( projHeme5, file = out_file )

    #plotPDF(p, name = "plotTrajectory.atac.pdf", ArchRProj = projHeme5, addDOC = FALSE, width = 5, height = 5)

    ###########################################################################################
    ## 展示基因的表达的时序变化
    ## Foxo1和Dmrt1是在未分化精原细胞（SSC）里面重要的
    ## E2f4是在分化的精原细胞里面活性较高的
    ## 早期还有Sohlh2，WT1
    gene_use <- c("FOXO1" , "DMRT1", "E2F4" , "SOHLH2" , "WT1")

    for( gene in gene_use ){
        ## 表达
        p <- plotTrajectory(projHeme5, trajectory = "Trajectory_time", colorBy = "GeneExpressionMatrix", name = gene, continuousSet = "blueYellow")
        image_name <- paste0( out_path , "/plotTrajectory." , gene , ".GeneExpressionMatrix." , level_name , ".pdf")
        pdf(image_name , width = 10, height = 8)
        print(p)
        dev.off()

        ## atac
        p <- plotTrajectory(projHeme5, trajectory = "Trajectory_time", colorBy = "GeneScoreMatrix", name = gene, continuousSet = "blueYellow")
        image_name <- paste0( out_path , "/plotTrajectory." , gene , ".GeneScoreMatrix." , level_name , ".pdf")
        pdf(image_name , width = 10, height = 8)
        print(p)
        dev.off()
    }

    ###########################################################################################
    ## 伪时间热图
    ## 通过整合基因评分/基因表达与跨伪时间的基序可访问性来识别正 TF 调节因子。

    ## atac活性程度
    trajGSM <- getTrajectory(ArchRProj = projHeme5, name = "Trajectory_time", useMatrix = "GeneScoreMatrix", log2Norm = TRUE)
    ## dim: 24919 100 

    ## 基因表达矩阵
    trajGEM <- getTrajectory(ArchRProj = projHeme5, name = "Trajectory_time", useMatrix = "GeneExpressionMatrix", log2Norm = FALSE)
    ## dim: 36438 100 

    ## motif活性
    trajMM  <- getTrajectory(ArchRProj = projHeme5, name = "Trajectory_time", useMatrix = "MotifMatrix", log2Norm = FALSE)
    ## dim: 1740 100 

    ## 去除基因名上的chr信息
    rownames(trajGSM) <- sapply(strsplit( rownames(trajGSM) , ":" ) , "[" , 2)
    rownames(trajGEM) <- sapply(strsplit( rownames(trajGEM) , ":" ) , "[" , 2)
    rownames(trajMM) <- sapply(strsplit( rownames(trajMM) , ":" ) , "[" , 2)

    ## 画图函数
    plotTrajectoryHeatmap_use <- function( trajGSM = trajGSM , trajGEM = trajGEM , image_name = image_name ){
        ## Integrative pseudo-time analyses
        ## 鉴定在时序过程中，RNA和ATAC改变趋势的基因
        corGSM_MM <- correlateTrajectories( trajGSM, trajGEM )
        ## 去除基因名上的chr信息
        corGSM_MM[[1]]$name1 <- sapply(strsplit( corGSM_MM[[1]]$name1 , ":" ) , "[" , 2)
        corGSM_MM[[1]]$name2 <- sapply(strsplit( corGSM_MM[[1]]$name2 , ":" ) , "[" , 2)

        ## 去除-AS1，反义转录本的基因
        use_gene <- grep( "-" , corGSM_MM[[1]]$name1 , invert = T)
        ## 去除不可靠的基因
        index <- which(corGSM_MM[[1]]$name1 %in% del_gene)
        if(length(index) > 0){
            use_gene <- use_gene[!use_gene %in% index]
        }

        ## 提取对应的基因
        trajGSM2 <- trajGSM[corGSM_MM[[1]]$name1[use_gene], ]
        trajGEM2 <- trajGEM[corGSM_MM[[1]]$name2[use_gene], ]

        ## 提取感兴趣的基因
        #gene_use <- c("FOXO1" , "DMRT1", "E2F4" , "SOHLH2" , "WT1")
        #trajGSM2 <- trajGSM[gene_use, ]
        #trajMM2 <- trajMM[gene_use, ]

        trajCombined <- trajGSM2
        assay(trajCombined , withDimnames=FALSE) <- t(apply(assay(trajGSM2), 1, scale)) + t(apply(assay(trajGEM2), 1, scale))
        combinedMat <- plotTrajectoryHeatmap(trajCombined, returnMat = TRUE, varCutOff = 0)

        rowOrder <- match(rownames(combinedMat), rownames(trajGSM2))
        ht1 <- plotTrajectoryHeatmap(trajGSM2,  pal = paletteContinuous(set = "horizonExtra"),  varCutOff = 0, rowOrder = rowOrder)
        ht2 <- plotTrajectoryHeatmap(trajGEM2, pal = paletteContinuous(set = "solarExtra"), varCutOff = 0, rowOrder = rowOrder)

        gene_list <- data.frame(gene = rownames(trajGSM2))
        
        write.table( gene_list , image_name , row.names = F , sep = "\t" , quote = F )

        return(ht1 + ht2)

    }

    ## 表达和atac评分
    image_name <- paste0( out_path , "/plotTrajectoryHeatmap.GEM_GSM.tsv")
    p <- plotTrajectoryHeatmap_use( trajGSM = trajGSM , trajGEM = trajGEM , image_name = image_name )
    ## 2023-12-26 17:29:20.725078 : Found 896 Correlated Pairings!
    image_name <- paste0( out_path , "/plotTrajectoryHeatmap.GEM_GSM.pdf")
    pdf(image_name , width = 10, height = 8)
    print(p)
    dev.off()

    ## 表达和motif评分
    image_name <- paste0( out_path , "/plotTrajectoryHeatmap.GEM_MM.tsv")
    p <- plotTrajectoryHeatmap_use( trajGSM = trajGEM , trajGEM = trajMM , image_name = image_name )
    ## 2023-12-26 17:30:50.674348 : Found 29 Correlated Pairings!
    image_name <- paste0( out_path , "/plotTrajectoryHeatmap.GEM_MM.pdf")
    pdf(image_name , width = 10, height = 8)
    print(p)
    dev.off()

    ## atac评分和motif评分
    image_name <- paste0( out_path , "/plotTrajectoryHeatmap.GSM_MM.tsv")
    p <- plotTrajectoryHeatmap_use( trajGSM = trajGSM , trajGEM = trajMM , image_name = image_name )
    ## 2023-12-26 17:32:00.304176 : Found 34 Correlated Pairings!
    image_name <- paste0( out_path , "/plotTrajectoryHeatmap.GSM_MM.pdf")
    pdf(image_name , width = 10, height = 8)
    print(p)
    dev.off()

}